#include <jni.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <unordered_map>

#include "safe_keyboard_lib.h"
#include <openssl/hmac.h>
#include <openssl/md5.h>

#include <jni.h>
#include <string.h>

using namespace std;

unordered_map<string, string> hashmap;

// AES加密密钥，生成规则是先随机生成一个密钥，然后根据生成的密钥再次进行MD5
static std::string aesSecretKey;

// 加密的iv，随机生成
static std::string aesSecretIV;

// 假设env是JNIEnv指针，obj是一个本地对象引用，str是要转换的std::string
jbyteArray stringToJByteArray(JNIEnv *env, const std::string &str) {
    // 创建一个新的jbyteArray
    jbyteArray jArray = env->NewByteArray(str.size());
    // 获取C风格的字符数组
    const jbyte *bytes = reinterpret_cast<const jbyte *>(str.c_str());
    // 复制字节到Java字节数组
    env->SetByteArrayRegion(jArray, 0, str.size(), bytes);
    return jArray;
}

string jstring2str(JNIEnv *env, jstring jstr) {
    char *rtn = NULL;
    jclass clsstring = env->FindClass("java/lang/String");
    jstring strencode = env->NewStringUTF("GB2312");
    jmethodID mid = env->GetMethodID(clsstring, "getBytes", "(Ljava/lang/String;)[B");
    jbyteArray barr = (jbyteArray) env->CallObjectMethod(jstr, mid, strencode);
    jsize alen = env->GetArrayLength(barr);
    jbyte *ba = env->GetByteArrayElements(barr, JNI_FALSE);
    if (alen > 0) {
        rtn = (char *) malloc(alen + 1);
        memcpy(rtn, ba, alen);
        rtn[alen] = 0;
    }
    env->ReleaseByteArrayElements(barr, ba, 0);
    std::string stemp(rtn);
    free(rtn);
    return stemp;
}

extern "C"
JNIEXPORT jbyteArray JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_encryption(JNIEnv *env, jclass type,
                                                   jbyteArray src_) {
    jbyte *src = env->GetByteArrayElements(src_, NULL);
    jsize src_Len = env->GetArrayLength(src_);
    int outlen = 0, cipherText_len;
    unsigned char *out = (unsigned char *) malloc((src_Len / 16 + 1) * 16);
    // 清空内存空间
    memset(out, 0, (src_Len / 16 + 1) * 16);
    EVP_CIPHER_CTX ctx;
    EVP_CIPHER_CTX_init(&ctx);
    // AES->指定加密算法，初始化加密key/iv
    // 这里可以修改签名算法：EVP_aes_128_cbc/EVP_aes_128_ecb/EVP_aes_128_cfb1/EVP_aes_128_cfb8
    EVP_EncryptInit_ex(&ctx, EVP_aes_128_cbc(), NULL, (const unsigned char *) aesSecretKey.c_str(),
                       (const unsigned char *) aesSecretIV.c_str());
    // AES->对数据进行加密运算
    EVP_EncryptUpdate(&ctx, out, &outlen, (const unsigned char *) src, src_Len);
    cipherText_len = outlen;
    // AES->结束加密运算
    EVP_EncryptFinal_ex(&ctx, out + outlen, &outlen);
    cipherText_len += outlen;
    // AES->EVP_CIPHER_CTX_cleanup
    EVP_CIPHER_CTX_cleanup(&ctx);
    // AES->从jni释放数据指针
    env->ReleaseByteArrayElements(src_, src, 0);
    jbyteArray cipher = env->NewByteArray(cipherText_len);
    // AES->在堆中分配ByteArray数组对象成功，将拷贝数据到数组中
    env->SetByteArrayRegion(cipher, 0, cipherText_len, (jbyte *) out);
    // AES->释放内存"
    free(out);
    return cipher;
}

extern "C"
JNIEXPORT jbyteArray JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_cryption(JNIEnv *env, jclass type, jbyteArray src_) {
    jbyte *src = env->GetByteArrayElements(src_, NULL);
    if (src == NULL) {
        return env->NewByteArray(0);
    }
    jsize src_Len = env->GetArrayLength(src_);
    int outlen = 0, plaintext_len;
    unsigned char *out = (unsigned char *) malloc(src_Len);
    memset(out, 0, src_Len);
    EVP_CIPHER_CTX ctx;
    EVP_CIPHER_CTX_init(&ctx);
    // AES->指定解密算法，初始化解密key/iv
    EVP_DecryptInit_ex(&ctx, EVP_aes_128_cbc(), NULL, (const unsigned char *) aesSecretKey.c_str(),
                       (const unsigned char *) aesSecretIV.c_str());
    // AES->对数据进行解密运算
    EVP_DecryptUpdate(&ctx, out, &outlen, (const unsigned char *) src, src_Len);
    plaintext_len = outlen;
    // AES->结束解密运算
    EVP_DecryptFinal_ex(&ctx, out + outlen, &outlen);
    plaintext_len += outlen;
    // AES->EVP_CIPHER_CTX_cleanup
    EVP_CIPHER_CTX_cleanup(&ctx);
    // AES->从jni释放数据指针
    env->ReleaseByteArrayElements(src_, src, 0);
    jbyteArray cipher = env->NewByteArray(plaintext_len);
    // AES->在堆中分配ByteArray数组对象成功，将拷贝数据到数组中
    env->SetByteArrayRegion(cipher, 0, plaintext_len, (jbyte *) out);
    // AES->释放内存
    free(out);
    return cipher;
}

std::string md5(const std::string &input) {
    unsigned char digest[16];
    if (input.empty()) {
        return "";
    }
    MD5((unsigned char *) input.data(), input.length(), digest);
    std::string result;
    for (int i = 0; i < sizeof(digest) / sizeof(digest[0]); ++i) {
        char hexDigit[(2 * sizeof(char)) + 1] = {'\0'};

        sprintf(hexDigit, "%02x", static_cast<unsigned int>(digest[i]));
        result += hexDigit;
    }
    return result;
}

// 生成随机密钥
string generateSecretKey() {
    // 初始化为16位长度的空字符串
    std::string Password(16, '\0');
    // 设置随机数种子
    srand(static_cast<unsigned>(time(nullptr)));
    // 随机密钥包含的内容
    const char SECRET_CHARS[] = "01234abcdefghijklMNOPQRSTUVWXYZ_";
    for (int i = 0; i < 16; ++i) {
        int random = rand() % strlen(SECRET_CHARS);
        Password[i] = SECRET_CHARS[random];
    }
    return Password;
}

// 生成随机向量
string generateSecretIV() {
    std::string Password(16, '\0'); // 初始化为16位长度的空字符串
    srand(static_cast<unsigned>(time(nullptr))); // 设置随机数种子
    // 随机密钥包含的内容
    const char IV_CHARS[] = "56789mnopqrstuvwxyzABCDEFGHIJKL_";
    for (int i = 0; i < 16; ++i) {
        int random = rand() % strlen(IV_CHARS);
        Password[i] = IV_CHARS[random];
    }
    return Password;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_generateSecretKey(JNIEnv *env, jclass type) {
    // 生成随机密钥和向量，保存到全局变量
    aesSecretKey = generateSecretKey();
    LOGD("########## generateSecretKey aesSecretKey = %s", aesSecretKey.c_str());
    aesSecretIV = generateSecretIV();
    LOGD("########## generateSecretKey aesSecretIV = %s", aesSecretIV.c_str());
}

extern "C"
JNIEXPORT void JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_addKey(JNIEnv *env, jclass type, jstring id_,
                                               jstring text_) {
    string input = jstring2str(env, text_);
    LOGD("########## addKey input = %s", input.c_str());
    hashmap[jstring2str(env, id_)] =
            hashmap[jstring2str(env, id_)] + input;
}

extern "C"
JNIEXPORT jstring JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_getKey(JNIEnv *env, jclass type, jstring id_) {
    string data = hashmap[jstring2str(env, id_)];
    if (data.c_str() == NULL) {
        // 没有加密数据，返回空的字符串
        return env->NewStringUTF("");
    }
    return env->NewStringUTF(data.c_str());
}


extern "C"
JNIEXPORT void JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_deleteKey(JNIEnv *env, jclass type, jstring id_) {
    string input = hashmap[jstring2str(env, id_)];
    input.pop_back();
    hashmap[jstring2str(env, id_)] = input;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_clearKey(JNIEnv *env, jclass type, jstring id_) {
    string input = hashmap[jstring2str(env, id_)];
    LOGD("########## clearKey input = %s", input.c_str());
    input.clear();
    hashmap[jstring2str(env, id_)] = input;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_surfin_safe_keyboard_JNIMethod_releaseKey(JNIEnv *env, jclass type, jstring id_) {
    const char *id = env->GetStringUTFChars(id_, 0);
    hashmap.erase(id);
    env->ReleaseStringUTFChars(id_, id);
}